#include <gmssl/error.h>
#include <gmssl/sm2.h>
unsigned char private_1[32] = {0x39, 0x45, 0x20, 0x8f, 0x7b, 0x21, 0x44, 0xb1,
                               0x3f, 0x36, 0xe3, 0x8a, 0xc6, 0xd3, 0x9f, 0x95,
                               0x88, 0x93, 0x93, 0x69, 0x28, 0x60, 0xb5, 0x1a,
                               0x42, 0xfb, 0x81, 0xef, 0x4d, 0xf7, 0xc5, 0xb8};

unsigned char public_1[64] = {
    0x09, 0xF9, 0xDF, 0x31, 0x1E, 0x54, 0x21, 0xA1, 0x50, 0xDD, 0x7D,
    0x16, 0x1E, 0x4B, 0xC5, 0xC6, 0x72, 0x17, 0x9F, 0xAD, 0x18, 0x33,
    0xFC, 0x07, 0x6B, 0xB0, 0x8F, 0xF3, 0x56, 0xF3, 0x50, 0x20, 0xCC,
    0xEA, 0x49, 0x0C, 0xE2, 0x67, 0x75, 0xA5, 0x2D, 0xC6, 0xEA, 0x71,
    0x8C, 0xC1, 0xAA, 0x60, 0x0A, 0xED, 0x05, 0xFB, 0xF3, 0x5E, 0x08,
    0x4A, 0x66, 0x32, 0xF6, 0x07, 0x2D, 0xA9, 0xAD, 0x13,
};

typedef struct SM2_KAP_CTX{
    SM2_KEY key;
    SM2_POINT rpub;
    SM2_POINT R;
    SM2_BN two_pow_w;
    SM2_BN t;

    unsigned char z[SM3_DIGEST_SIZE];
    unsigned char rz[SM3_DIGEST_SIZE];
    unsigned char sum[SM3_DIGEST_SIZE];
    int do_checksum;
    int is_initiator;
}SM2_KAP_CTX;

void print_bn(const char* str, SM2_BN bn)
{
    printf("%s: ", str);
    for (int i = 0; i < 8; i++) {
        printf("%08X", bn[i] & 0xFFFFFFFF);
        printf(" ");
    }
    printf("\n");
}

int SM2_KAP_CTX_init(SM2_KAP_CTX* ctx, SM2_KEY* key, SM2_POINT* rpub,
                     const char* id, size_t idlen, const char* rid,
                     size_t ridlen, const int is_initiator, const int do_checksum);
int SM2_KAP_prepare(SM2_KAP_CTX* ctx, unsigned char R[64]);
int SM2_KAP_compute_key(SM2_KAP_CTX* ctx, unsigned char R[64],
                        unsigned char* key, size_t keylen,
                        unsigned char checksum[SM3_DIGEST_SIZE]);
int SM2_KAP_final_check(SM2_KAP_CTX*        ctx,
                        const unsigned char checksum[SM3_DIGEST_SIZE]);
void SM2_KAP_CTX_cleanup(SM2_KAP_CTX* ctx);
// 0x80000000 00000000 00000000 00000000
// 00000000 00000000 00000000 00000000 80000000 00000000 00000000 00000000
#define KEYLEN 16
int main(int argc, char const *argv[])
{
    int         rc  = 0;
    const char* Aid = "123456";
    const char* Bid = "123456";

    unsigned char Ra[64];
    unsigned char Rb[64];
    unsigned char kab[KEYLEN];
    unsigned char kba[KEYLEN];
    unsigned char suma[SM3_DIGEST_SIZE];
    unsigned char sumb[SM3_DIGEST_SIZE];

    SM2_KEY keyA;
    SM2_KEY keyB;
    sm2_key_generate(&keyA);
    sm2_key_generate(&keyB);
    // memcpy(keyA.private_key, private_1, 32);
    // memcpy(&keyA.public_key, public_1, 64);
    // memcpy(keyB.private_key, private_1, 32);
    // memcpy(&keyB.public_key, public_1, 64);

    SM2_KAP_CTX ctxA;
    SM2_KAP_CTX ctxB;

    rc = SM2_KAP_CTX_init(&ctxA, &keyA, &keyB.public_key, Aid, strlen(Aid), Bid,
                          strlen(Bid), 1, 0);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }
    rc = SM2_KAP_CTX_init(&ctxB, &keyB, &keyA.public_key, Bid, strlen(Bid), Aid,
                          strlen(Aid), 0, 0);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }

    rc = SM2_KAP_prepare(&ctxA, Ra);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }
    rc = SM2_KAP_prepare(&ctxB, Rb);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }

    rc = SM2_KAP_compute_key(&ctxA, Rb, kab, KEYLEN, suma);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }
    rc = SM2_KAP_compute_key(&ctxB, Ra, kba, KEYLEN, sumb);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
        goto end;
    }

    rc = SM2_KAP_final_check(&ctxA, sumb);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
    }
    rc = SM2_KAP_final_check(&ctxB, suma);
    if (rc) {
        printf("[%d] error rc = %d\n", __LINE__, rc);
    }

    printf("key:\n");
    print_bytes(kab, KEYLEN);
    print_bytes(kba, KEYLEN);
end:
    SM2_KAP_CTX_cleanup(&ctxA);
    SM2_KAP_CTX_cleanup(&ctxB);
    return 0;
}

SM2_BN two_pow_w = {0x0, 0x0, 0x0, 0x80000000, 0x0, 0x0, 0x0, 0x0};
SM2_BN SM2_H = {1,0,0,0,0,0,0,0};

int SM2_KAP_CTX_init(SM2_KAP_CTX* ctx, SM2_KEY* key, SM2_POINT* rpub,
                     const char* id, size_t idlen, const char* rid,
                     size_t ridlen, const int is_initiator, const int do_checksum)
{
    ctx->do_checksum = do_checksum;
    ctx->is_initiator = is_initiator;
    sm2_bn_set_zero(ctx->t);
    memset(&ctx->R, 0, sizeof(ctx->R));

    memcpy(&ctx->key, key, sizeof(ctx->key));
    memcpy(&ctx->rpub, rpub, sizeof(ctx->rpub));


    sm2_compute_z(ctx->z, &ctx->key.public_key, id, idlen);
    sm2_compute_z(ctx->rz, rpub, rid, ridlen);

    return 0;
}

void SM2_KAP_CTX_cleanup(SM2_KAP_CTX* ctx)
{
    memset(ctx, 0 ,sizeof(*ctx));
}

void sm2_bn_and(SM2_BN ret, const SM2_BN a, const SM2_BN b)
{
    for (int i = 0; i < 8; i++) {
        ret[i] = a[i] & b[i];
    }
}

unsigned char randm[]={
    0x43, 0x3D, 0x60, 0x93, 0x0D, 0xCE, 0x3B, 0x58,
    0xFD, 0xB3, 0xB3, 0x4E, 0x83, 0xA0, 0xEB, 0xFA,
    0xA5, 0x6A, 0x09, 0x44, 0xBF, 0xC3, 0xCE, 0x3F,
    0x8E, 0xA6, 0x21, 0x13, 0xF2, 0x84, 0xFD, 0xB9};

extern SM2_BN SM2_ONE;
extern SM2_BN SM2_N;

int SM2_KAP_prepare(SM2_KAP_CTX* ctx, unsigned char R[64])
{
    SM2_BN rA;
    SM2_BN _x;
    SM2_BN x,y;
    SM2_JACOBIAN_POINT _rG, *rG = &_rG;

    sm2_bn_set_zero(x);
    sm2_bn_set_zero(rA);
    sm2_bn_set_zero(y);
    sm2_bn_set_zero(_x);
    sm2_jacobian_point_init(rG);

    /* 第一步 产生随机数 rA∈[1,n-1]; */
	do {
		sm2_bn_rand_range(rA, SM2_N);
	} while (sm2_bn_is_zero(rA));
    // sm2_bn_from_bytes(rA, randm);
  
    /* 第二步 生成rG 计算椭圆曲线点RA= [rA]G=(x1,y1); */
    sm2_jacobian_point_mul_generator(rG, rA);
    sm2_jacobian_point_get_xy(rG, x, y);
    sm2_bn_to_bytes(x, ctx->R.x);
    sm2_bn_to_bytes(y, ctx->R.y);

    /* 第四步 从RA中取出域元素x1 计算
        x_ = 2^w + (x1 & (2^w -1)) */
    sm2_bn_sub(_x, two_pow_w, SM2_ONE);
    sm2_bn_and(_x, x, _x);
    sm2_bn_add(_x, two_pow_w, _x);

    /* 第五步 计算 t=(d + x_·rA)mod n; */
    SM2_BN d;
    sm2_bn_from_bytes(d, ctx->key.private_key);
    sm2_fn_mul(ctx->t, _x, rA); // different from mpi
    sm2_fn_add(ctx->t, d, ctx->t);

    /* 第三步 输出R 将RA发送给用户B; */
    memcpy(R, &ctx->R, 64);
    return 0;
}

extern int sm2_kdf(const uint8_t *in, size_t inlen, size_t outlen, uint8_t *out);

static void SM3_S(uint8_t flage, uint8_t U[64], uint8_t Z[32], uint8_t rZ[32],
    uint8_t p1[64], uint8_t p2[64], uint8_t S[32]);

int SM2_KAP_compute_key(SM2_KAP_CTX* ctx, unsigned char R[64],
                        unsigned char* key, size_t keylen,
                        unsigned char checksum[SM3_DIGEST_SIZE])
{
    int rc = 0;
    SM2_BN x,y, _x;
    SM2_JACOBIAN_POINT _Rp, *Rp = &_Rp, _U, *U = &_U;

    sm2_bn_set_zero(x);
    sm2_bn_set_zero(y);
    sm2_bn_set_zero(_x);
    sm2_jacobian_point_init(Rp);
    sm2_jacobian_point_init(U);

    /* 第六步：验证RB是否满足椭圆曲线方程 */
    sm2_jacobian_point_from_bytes(Rp, R);
    rc = sm2_jacobian_point_is_on_curve(Rp);
    if(rc < 0)
        printf("sm2 not in curve\n");

    /* 第六步：从RB中取出域元素x2, 计算
        x_ = 2^w + (x2 and (2^w - 1)) */
    sm2_bn_sub(_x, two_pow_w, SM2_ONE);
    sm2_bn_and(_x, Rp->X, _x);
    sm2_bn_add(_x, two_pow_w, _x);

    /* 第七步：U = ht * (P + x_ * R)=(xU ,yU ), check U != O */
    SM2_JACOBIAN_POINT p;
    sm2_jacobian_point_from_bytes(&p, (const uint8_t*)&ctx->rpub);

    sm2_jacobian_point_mul(U, _x, Rp);
    sm2_jacobian_point_get_xy(U, x, y);
    sm2_jacobian_point_set_xy(U, x, y);
    sm2_jacobian_point_add(U, &p, U);
    sm2_jacobian_point_mul(U, ctx->t, U);
    sm2_jacobian_point_mul(U, SM2_H, U);
    sm2_jacobian_point_get_xy(U, x, y);
    sm2_jacobian_point_set_xy(U, x, y);
    if (sm2_bn_is_zero(U->Z))
        return -1;
    sm2_jacobian_point_get_xy(U, x, y);

    unsigned char kdfbuf[128] = {0};
    sm2_bn_to_bytes(x, kdfbuf);
    sm2_bn_to_bytes(y, kdfbuf+32);
    if (ctx->is_initiator) {
        memcpy(kdfbuf + 64, ctx->z, SM3_DIGEST_SIZE);
        memcpy(kdfbuf + 64 + SM3_DIGEST_SIZE, ctx->rz,
               SM3_DIGEST_SIZE);
    } else {
        memcpy(kdfbuf + 64, ctx->rz, SM3_DIGEST_SIZE);
        memcpy(kdfbuf + 64 + SM3_DIGEST_SIZE, ctx->z,
               SM3_DIGEST_SIZE);
    }
    sm2_kdf(kdfbuf, 128, keylen, key);

    /* 第九步（可选）：SM3(0x02‖yU‖SM3(xU‖ZA‖ZB‖x1‖y1‖x2‖y2)
       第十步（可选）：SM3(0x03‖yU‖SM3(xU‖ZA‖ZB‖x1‖y1‖x2‖y2) */
    if (ctx->do_checksum && checksum) {
        if (ctx->is_initiator) {
            SM3_S(0x02, kdfbuf, ctx->z, ctx->rz, (uint8_t*)&ctx->R, R, ctx->sum);
            SM3_S(0x03, kdfbuf, ctx->z, ctx->rz, (uint8_t*)&ctx->R, R, checksum);
        } else {
            SM3_S(0x03, kdfbuf, ctx->rz, ctx->z, R, (uint8_t*)&ctx->R, ctx->sum);
            SM3_S(0x02, kdfbuf, ctx->rz, ctx->z, R, (uint8_t*)&ctx->R, checksum);
        }
    }
    return 0;
}

int SM2_KAP_final_check(SM2_KAP_CTX*        ctx,
                        const unsigned char checksum[SM3_DIGEST_SIZE])
{
    if (!ctx || !checksum) return -1;
    if (ctx->do_checksum) {
        if (memcmp(ctx->sum, checksum, SM3_DIGEST_SIZE)) {
            return 1;
        }
    }
    return 0;
}

static void SM3_S(uint8_t flage, uint8_t U[64], uint8_t Z[32], uint8_t rZ[32],
    uint8_t p1[64], uint8_t p2[64], uint8_t S[32])
{
    SM3_CTX ctx[1];
    uint8_t   tmp[32];
    sm3_init(ctx);

    sm3_update(ctx, U, 32);
    sm3_update(ctx, Z, 32);
    sm3_update(ctx, p1, 32);
    sm3_update(ctx, p1 + 32, 32);
    sm3_update(ctx, p2, 32);
    sm3_update(ctx, p2 + 32, 32);
    sm3_finish(ctx, tmp);

    sm3_init(ctx);
    sm3_update(ctx, &flage, 1);
    sm3_update(ctx, U + 32, 32);
    sm3_update(ctx, tmp, 32);
    sm3_finish(ctx, S);
}